"""初始化预定义角色和权限"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import inspect, text
from sqlalchemy.orm import Session
from database.db import engine, Base, get_db
from models.user import Role, Permission

def add_level_column():
    """为roles表添加level列"""
    inspector = inspect(engine)
    columns = [col['name'] for col in inspector.get_columns('roles')]
    
    if 'level' not in columns:
        with engine.connect() as conn:
            conn.execute(text('ALTER TABLE roles ADD COLUMN level INTEGER DEFAULT 0 NOT NULL'))
            conn.commit()
            print("已成功为roles表添加level列")
    else:
        print("roles表已包含level列")

def init_roles_and_permissions(db: Session):
    """初始化角色和权限"""
    
    # 定义预定义角色
    predefined_roles = [
        {"name": "超级管理员", "description": "系统配置管理员", "level": 500},
        {"name": "任务管理员", "description": "创建和分配无人机任务", "level": 400},
        {"name": "飞手", "description": "操作无人机执行任务", "level": 300},
        {"name": "观测员", "description": "查看无人机画面和数据", "level": 200},
        {"name": "审计员", "description": "查看系统操作日志", "level": 100}
    ]
    
    # 定义权限
    permissions = [
        {"name": "system_config", "description": "系统配置权限", "roles": ["超级管理员"]},
        {"name": "user_manage", "description": "用户管理权限", "roles": ["超级管理员"]},
        {"name": "role_manage", "description": "角色管理权限", "roles": ["超级管理员"]},
        {"name": "permission_manage", "description": "权限管理权限", "roles": ["超级管理员"]},
        {"name": "department_manage", "description": "部门管理权限", "roles": ["超级管理员", "任务管理员"]},
        {"name": "drone_manage", "description": "无人机管理权限", "roles": ["超级管理员", "任务管理员"]},
        {"name": "task_create", "description": "创建任务权限", "roles": ["任务管理员"]},
        {"name": "task_assign", "description": "分配任务权限", "roles": ["任务管理员"]},
        {"name": "drone_control", "description": "控制无人机权限", "roles": ["飞手"]},
        {"name": "flight_view", "description": "查看飞行画面权限", "roles": ["飞手", "观测员", "任务管理员"]},
        {"name": "data_view", "description": "查看数据权限", "roles": ["飞手", "观测员", "任务管理员", "审计员"]},
        {"name": "log_view", "description": "查看日志权限", "roles": ["审计员", "超级管理员"]},
        {"name": "route_manage", "description": "航线管理权限", "roles": ["任务管理员", "超级管理员"]},
    ]
    
    # 先创建permissions表（如果不存在）
    if not inspect(engine).has_table('permissions'):
        Base.metadata.create_all(bind=engine, tables=[Permission.__table__])
        print("已创建permissions表")
    
    # 创建权限
    permission_dict = {}
    for perm_data in permissions:
        permission = db.query(Permission).filter(Permission.name == perm_data["name"]).first()
        if not permission:
            permission = Permission(
                name=perm_data["name"],
                description=perm_data["description"]
            )
            db.add(permission)
            db.commit()
            db.refresh(permission)
        permission_dict[perm_data["name"]] = permission
    
    # 创建或更新角色
    for role_data in predefined_roles:
        # 使用原生SQL查询角色，避免ORM层面的字段检查问题
        result = db.execute(text("SELECT id FROM roles WHERE name = :name"), {"name": role_data["name"]})
        role_id = result.scalar()
        
        if role_id:
            # 更新现有角色
            db.execute(
                text("UPDATE roles SET description = :description, level = :level WHERE id = :id"),
                {
                    "description": role_data["description"],
                    "level": role_data["level"],
                    "id": role_id
                }
            )
            db.commit()
            role = db.query(Role).filter(Role.id == role_id).first()
            print(f"已更新角色: {role_data['name']}, 级别: {role_data['level']}")
        else:
            # 创建新角色
            try:
                role = Role(
                    name=role_data["name"],
                    description=role_data["description"],
                    level=role_data["level"]
                )
                db.add(role)
                db.commit()
                db.refresh(role)
                print(f"已创建角色: {role_data['name']}, 级别: {role_data['level']}")
            except Exception as e:
                print(f"创建角色 {role_data['name']} 失败: {str(e)}")
                db.rollback()
                continue
        
        # 为角色分配权限（先清除现有权限）
        if role:
            try:
                # 清除现有权限关联
                db.execute(
                    text("DELETE FROM role_permissions WHERE role_id = :role_id"),
                    {"role_id": role.id}
                )
                
                # 添加新权限关联
                for perm_data in permissions:
                    if role_data["name"] in perm_data["roles"]:
                        permission = permission_dict.get(perm_data["name"])
                        if permission:
                            db.execute(
                                text("INSERT OR IGNORE INTO role_permissions (role_id, permission_id) VALUES (:role_id, :permission_id)"),
                                {"role_id": role.id, "permission_id": permission.id}
                            )
                
                db.commit()
                # 重新加载角色以获取最新的权限
                db.refresh(role)
                print(f"已为角色 {role.name} 分配权限: {len(role.permissions)} 个")
            except Exception as e:
                print(f"为角色 {role.name} 分配权限失败: {str(e)}")
                db.rollback()

if __name__ == "__main__":
    try:
        # 先添加level列（如果不存在）
        add_level_column()
        
        # 获取数据库会话
        db = next(get_db())
        
        try:
            init_roles_and_permissions(db)
            print("\n角色和权限初始化完成！")
        except Exception as e:
            print(f"\n初始化失败: {str(e)}")
            db.rollback()
        finally:
            db.close()
    except Exception as e:
        print(f"数据库操作失败: {str(e)}")